from torchvision.datasets import VOCDetection
from torch.utils.data import DataLoader
from utils import collate_fn_voc
from transforms import RandomFlip, Resize, ToTensor, Normalize, Compose
from torchvision.ops import box_convert
import torch

class VocDataset(VOCDetection):

    def __init__(self, root, year, image_set, voc_transforms=None):
        super().__init__(root=root, year=year, image_set=image_set)
        self.voc_transforms = voc_transforms

    def __getitem__(self, idx):
        img, target = super().__getitem__(idx)
        w, h = img.size
        img_id = target['annotation']['filename'][:6]
        labels = [voc2012_class_to_idx[item['name']] for item in target['annotation']['object']]
        boxes = [[int(item['bndbox']['xmin']), int(item['bndbox']['ymin']), int(item['bndbox']['xmax']), int(item['bndbox']['ymax'])] \
                 for item in target['annotation']['object']]
        boxes = box_convert(torch.tensor(boxes), in_fmt='xyxy', out_fmt='xywh').tolist()
        if self.voc_transforms:
            img, boxes = self.voc_transforms(img, boxes)
        return img, (boxes, labels), img_id, (w, h)



voc2012_class_to_idx = {
    "aeroplane": 0,
    "bicycle": 1,
    "bird": 2,
    "boat": 3,
    "bottle": 4,
    "bus": 5,
    "car": 6,
    "cat": 7,
    "chair": 8,
    "cow": 9,
    "diningtable": 10,
    "dog": 11,
    "horse": 12,
    "motorbike": 13,
    "person": 14,
    "pottedplant": 15,
    "sheep": 16,
    "sofa": 17,
    "train": 18,
    "tvmonitor": 19
}

#   0: aeroplane
#   1: bicycle
#   2: bird
#   3: boat
#   4: bottle
#   5: bus
#   6: car
#   7: cat
#   8: chair
#   9: cow
#   10: diningtable
#   11: dog
#   12: horse
#   13: motorbike
#   14: person
#   15: pottedplant
#   16: sheep
#   17: sofa
#   18: train
#   19: tvmonitor



if __name__ == "__main__":
    root = '/mnt/sdb2/ray/rtdetr-implement/'
    train_transforms = Compose([Resize((500, 500)),
                            RandomFlip(),
                            ToTensor(),
                            Normalize()])
    voc_train_dataset = VocDataset(root=root, year="2012", image_set="trainval", voc_transforms=train_transforms)
    dataloader = DataLoader(voc_train_dataset, batch_size=1, collate_fn=collate_fn_voc)
    print(len(voc_train_dataset))

    